#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
文档处理器
Document Processor for Local File Support

支持读取和处理本地文档文件，包括文本文件、PDF、Word文档等
"""

import os
import re
import logging
from typing import List, Dict, Optional, Tuple
from pathlib import Path
import mimetypes

# 尝试导入可选依赖
try:
    import PyPDF2
    PDF_AVAILABLE = True
except ImportError:
    PDF_AVAILABLE = False

try:
    from docx import Document
    DOCX_AVAILABLE = True
except ImportError:
    DOCX_AVAILABLE = False

try:
    import pandas as pd
    EXCEL_AVAILABLE = True
except ImportError:
    EXCEL_AVAILABLE = False

logger = logging.getLogger(__name__)

class DocumentProcessor:
    """文档处理器"""
    
    def __init__(self):
        self.supported_extensions = {
            '.txt': self._read_text_file,
            '.md': self._read_text_file,
            '.csv': self._read_csv_file,
            '.json': self._read_json_file,
            '.xml': self._read_xml_file,
            '.html': self._read_html_file,
            '.htm': self._read_html_file,
        }
        
        # 添加可选支持的格式
        if PDF_AVAILABLE:
            self.supported_extensions['.pdf'] = self._read_pdf_file
        
        if DOCX_AVAILABLE:
            self.supported_extensions['.docx'] = self._read_docx_file
            self.supported_extensions['.doc'] = self._read_docx_file
        
        if EXCEL_AVAILABLE:
            self.supported_extensions['.xlsx'] = self._read_excel_file
            self.supported_extensions['.xls'] = self._read_excel_file
    
    def get_supported_formats(self) -> List[str]:
        """获取支持的文件格式"""
        return list(self.supported_extensions.keys())
    
    def can_process_file(self, file_path: str) -> bool:
        """检查是否可以处理指定文件"""
        if not os.path.exists(file_path):
            return False
        
        file_ext = Path(file_path).suffix.lower()
        return file_ext in self.supported_extensions
    
    def process_document(self, file_path: str) -> Dict:
        """
        处理文档文件
        
        Args:
            file_path: 文件路径
            
        Returns:
            Dict: 包含文档内容的字典
        """
        try:
            if not self.can_process_file(file_path):
                return {
                    'success': False,
                    'error': f'不支持的文件格式: {Path(file_path).suffix}',
                    'supported_formats': self.get_supported_formats()
                }
            
            file_ext = Path(file_path).suffix.lower()
            reader_func = self.supported_extensions[file_ext]
            
            content = reader_func(file_path)
            
            return {
                'success': True,
                'content': content,
                'file_path': file_path,
                'file_size': os.path.getsize(file_path),
                'file_type': file_ext,
                'word_count': len(content.split()) if content else 0
            }
            
        except Exception as e:
            logger.error(f"处理文档失败 {file_path}: {str(e)}")
            return {
                'success': False,
                'error': f'处理文档失败: {str(e)}'
            }
    
    def _read_text_file(self, file_path: str) -> str:
        """读取文本文件"""
        try:
            # 尝试不同的编码
            encodings = ['utf-8', 'gbk', 'gb2312', 'latin-1']
            
            for encoding in encodings:
                try:
                    with open(file_path, 'r', encoding=encoding) as f:
                        return f.read()
                except UnicodeDecodeError:
                    continue
            
            # 如果所有编码都失败，使用二进制模式读取
            with open(file_path, 'rb') as f:
                content = f.read()
                return content.decode('utf-8', errors='ignore')
                
        except Exception as e:
            logger.error(f"读取文本文件失败 {file_path}: {str(e)}")
            return ""
    
    def _read_csv_file(self, file_path: str) -> str:
        """读取CSV文件"""
        try:
            if EXCEL_AVAILABLE:
                df = pd.read_csv(file_path)
                return df.to_string()
            else:
                # 如果没有pandas，使用基本的文本读取
                return self._read_text_file(file_path)
        except Exception as e:
            logger.error(f"读取CSV文件失败 {file_path}: {str(e)}")
            return self._read_text_file(file_path)
    
    def _read_json_file(self, file_path: str) -> str:
        """读取JSON文件"""
        try:
            import json
            with open(file_path, 'r', encoding='utf-8') as f:
                data = json.load(f)
                return json.dumps(data, ensure_ascii=False, indent=2)
        except Exception as e:
            logger.error(f"读取JSON文件失败 {file_path}: {str(e)}")
            return self._read_text_file(file_path)
    
    def _read_xml_file(self, file_path: str) -> str:
        """读取XML文件"""
        try:
            import xml.etree.ElementTree as ET
            tree = ET.parse(file_path)
            root = tree.getroot()
            
            # 提取文本内容
            text_content = []
            for elem in root.iter():
                if elem.text and elem.text.strip():
                    text_content.append(elem.text.strip())
            
            return '\n'.join(text_content)
        except Exception as e:
            logger.error(f"读取XML文件失败 {file_path}: {str(e)}")
            return self._read_text_file(file_path)
    
    def _read_html_file(self, file_path: str) -> str:
        """读取HTML文件"""
        try:
            from bs4 import BeautifulSoup
            
            content = self._read_text_file(file_path)
            soup = BeautifulSoup(content, 'html.parser')
            
            # 移除脚本和样式标签
            for script in soup(["script", "style"]):
                script.decompose()
            
            # 提取文本
            text = soup.get_text()
            
            # 清理文本
            lines = (line.strip() for line in text.splitlines())
            chunks = (phrase.strip() for line in lines for phrase in line.split("  "))
            text = '\n'.join(chunk for chunk in chunks if chunk)
            
            return text
        except ImportError:
            # 如果没有BeautifulSoup，使用基本的文本读取
            return self._read_text_file(file_path)
        except Exception as e:
            logger.error(f"读取HTML文件失败 {file_path}: {str(e)}")
            return self._read_text_file(file_path)
    
    def _read_pdf_file(self, file_path: str) -> str:
        """读取PDF文件"""
        if not PDF_AVAILABLE:
            return "PDF处理需要安装PyPDF2: pip install PyPDF2"
        
        try:
            text_content = []
            
            with open(file_path, 'rb') as file:
                pdf_reader = PyPDF2.PdfReader(file)
                
                for page_num in range(len(pdf_reader.pages)):
                    page = pdf_reader.pages[page_num]
                    text_content.append(page.extract_text())
            
            return '\n'.join(text_content)
        except Exception as e:
            logger.error(f"读取PDF文件失败 {file_path}: {str(e)}")
            return f"PDF读取失败: {str(e)}"
    
    def _read_docx_file(self, file_path: str) -> str:
        """读取Word文档"""
        if not DOCX_AVAILABLE:
            return "Word文档处理需要安装python-docx: pip install python-docx"
        
        try:
            doc = Document(file_path)
            text_content = []
            
            for paragraph in doc.paragraphs:
                if paragraph.text.strip():
                    text_content.append(paragraph.text)
            
            return '\n'.join(text_content)
        except Exception as e:
            logger.error(f"读取Word文档失败 {file_path}: {str(e)}")
            return f"Word文档读取失败: {str(e)}"
    
    def _read_excel_file(self, file_path: str) -> str:
        """读取Excel文件"""
        if not EXCEL_AVAILABLE:
            return "Excel文件处理需要安装pandas: pip install pandas openpyxl"
        
        try:
            df = pd.read_excel(file_path)
            return df.to_string()
        except Exception as e:
            logger.error(f"读取Excel文件失败 {file_path}: {str(e)}")
            return f"Excel文件读取失败: {str(e)}"
    
    def extract_key_information(self, content: str) -> List[str]:
        """
        从文档内容中提取关键信息
        
        Args:
            content: 文档内容
            
        Returns:
            List[str]: 关键信息列表
        """
        if not content:
            return []
        
        # 提取关键信息模式
        patterns = [
            r'研究表明.*?[。！？]',
            r'数据显示.*?[。！？]',
            r'专家.*?表示.*?[。！？]',
            r'根据.*?报告.*?[。！？]',
            r'统计.*?显示.*?[。！？]',
            r'结论.*?[。！？]',
            r'发现.*?[。！？]',
            r'分析.*?[。！？]'
        ]
        
        key_info = []
        for pattern in patterns:
            matches = re.findall(pattern, content, re.IGNORECASE)
            key_info.extend(matches)
        
        # 去重并限制数量
        unique_info = list(set(key_info))
        return unique_info[:10]  # 最多返回10条关键信息
    
    def get_document_summary(self, content: str) -> Dict:
        """
        获取文档摘要信息
        
        Args:
            content: 文档内容
            
        Returns:
            Dict: 摘要信息
        """
        if not content:
            return {
                'word_count': 0,
                'paragraph_count': 0,
                'key_info_count': 0,
                'estimated_reading_time': 0
            }
        
        # 统计信息
        words = content.split()
        paragraphs = [p for p in content.split('\n') if p.strip()]
        key_info = self.extract_key_information(content)
        
        # 估算阅读时间（假设每分钟200字）
        reading_time = len(words) / 200
        
        return {
            'word_count': len(words),
            'paragraph_count': len(paragraphs),
            'key_info_count': len(key_info),
            'estimated_reading_time': round(reading_time, 1),
            'key_info_preview': key_info[:3]  # 前3条关键信息
        } 